from PIL import Image, ImageDraw, ImageFont
import os
from pathlib import Path
from typing import Union

class Watermarker:
    def __init__(self, config: dict):
        self.validate_config(config)
        self.config = config
        self.prepare_directories()

    def validate_config(self, config):
        required_keys = ['source_dir', 'output_dir', 'watermark_type']
        for key in required_keys:
            if key not in config:
                raise ValueError(f"缺少必要参数: {key}")

        if config['watermark_type'] == 'text':
            if 'text_config' not in config or 'font_size' not in config['text_config']:
                raise ValueError("文字水印需要配置字体参数")
        elif config['watermark_type'] == 'image':
            if not os.path.exists(config['watermark_content']):
                raise FileNotFoundError("水印图片不存在")

    def prepare_directories(self):
        Path(self.config['output_dir']).mkdir(parents=True, exist_ok=True)

    def process_all(self):
        supported_formats = ('.jpg', '.jpeg', '.png', '.webp')
        for file in Path(self.config['source_dir']).rglob('*'):
            if file.suffix.lower() in supported_formats:
                try:
                    self.process_image(file)
                except Exception as e:
                    print(f"处理失败 {file}: {str(e)}")

    def process_image(self, img_path: Union[str, Path]):
        with Image.open(img_path) as img:
            watermark = self.create_watermark(img.size)
            combined = Image.alpha_composite(
                img.convert('RGBA'),
                watermark
            ).convert('RGB')
            output_path = self.get_output_path(img_path)
            combined.save(output_path, quality=95)

    def create_watermark(self, img_size: tuple) -> Image.Image:
        if self.config['watermark_type'] == 'text':
            return self.create_text_watermark(img_size)
        else:
            return self.create_image_watermark(img_size)

    def create_text_watermark(self, img_size: tuple) -> Image.Image:
        font = ImageFont.truetype(
            self.config['text_config']['font_path'],
            self.config['text_config']['font_size']
        )

        text = self.config['watermark_content']
        text_bbox = font.getbbox(text)
        text_width = text_bbox[2] - text_bbox[0]
        text_height = text_bbox[3] - text_bbox[1]

        watermark = Image.new('RGBA', (text_width, text_height), (0, 0, 0, 0))
        draw = ImageDraw.Draw(watermark)
        
        # 修正透明度: 将 0-100 转换为 0-255
        opacity = int(self.config['opacity'] * 255 / 100)
        fill_color = self.config['text_config']['color'] + (opacity,)

        draw.text(
            (0, 0),
            text,
            font=font,
            fill=fill_color
        )
        
        angle = self.config.get('rotation', 0)
        watermark = watermark.rotate(angle, expand=True)  # 旋转水印
        
        position = self.calculate_position(img_size, watermark.size)
        final_watermark = Image.new('RGBA', img_size, (0, 0, 0, 0))
        final_watermark.paste(watermark, position, watermark)
        
        return final_watermark

    def create_image_watermark(self, img_size: tuple) -> Image.Image:
        with Image.open(self.config['watermark_content']) as wm_img:
            wm_img = wm_img.convert('RGBA')
            wm_img = self.resize_watermark(wm_img, img_size)
            position = self.calculate_position(img_size, wm_img.size)
            
            watermark = Image.new('RGBA', img_size, (0,0,0,0))
            watermark.paste(wm_img, position, wm_img)
            return watermark

    def resize_watermark(self, wm_img: Image.Image, img_size: tuple) -> Image.Image:
        max_size = min(img_size) // 4  # 水印最大边长为原图1/4
        ratio = min(max_size / max(wm_img.size), 1)
        return wm_img.resize(
            tuple(int(dim * ratio) for dim in wm_img.size),
            Image.Resampling.LANCZOS
        )

    def calculate_position(self, img_size: tuple, wm_size: tuple) -> tuple:
        width, height = img_size
        wm_width, wm_height = wm_size
        margin = self.config['margin']
        
        position_map = {
            ('left', 'top'): (margin, margin),
            ('right', 'top'): (width - wm_width - margin, margin),
            ('left', 'bottom'): (margin, height - wm_height - margin),
            ('right', 'bottom'): (width - wm_width - margin, height - wm_height - margin),
            ('center', 'center'): ((width - wm_width)//2, (height - wm_height)//2)
        }
        return position_map.get(tuple(self.config['position']), (0,0))

    def get_output_path(self, src_path: Path) -> Path:
        relative_path = src_path.relative_to(self.config['source_dir'])
        return Path(self.config['output_dir']) / relative_path

if __name__ == "__main__":
    config = {
        'source_dir': 'F:\\微信公众号\\ACDSee 2025.zip',
        'output_dir': 'F:\\微信公众号-影子\\ACDSee 2025.zip',
        'watermark_type': 'text',
        'watermark_content': '微信公众号：软件联盟 Software Union',
        'position': ('center', 'center'),
        'margin': 20,
        'opacity': 255,  # 调整为 0-255，255 表示完全不透明
        'text_config': {
            'font_path': 'C://Windows//Fonts//msyh.ttc',
            'font_size': 30,
            'color': (255, 0, 0)  # 纯红色
        },
        'rotation': -30
    }

    watermarker = Watermarker(config)
    watermarker.process_all()